#include "lua/lualib.h"
#include "lua/lauxlib.h"

#include <cassert>
#include <cstring>

#include "core.hpp"
#include "../../scripting/lua.hpp"
#include "../../scripting/lua_api.hpp"

#include "../../actions.hpp"
#include "../../attack_prediction.hpp"
#include "../../filesystem.hpp"
#include "../../foreach.hpp"
#include "../../game_display.hpp"
#include "../../gamestatus.hpp"
#include "../../log.hpp"
#include "../../map.hpp"
#include "../../pathfind/pathfind.hpp"
#include "../../play_controller.hpp"
#include "../../resources.hpp"
#include "../../terrain_translation.hpp"
#include "../../unit.hpp"
#include "../actions.hpp"
#include "../composite/engine_lua.hpp"

static lg::log_domain log_ai_engine_lua("ai/engine/lua");
#define LOG_LUA LOG_STREAM(info, log_ai_engine_lua)
#define ERR_LUA LOG_STREAM(err, log_ai_engine_lua)

static char const aisKey = 0;

namespace ai
{
	void lua_ai_context::init(lua_State* L)
	{
		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_newtable(L);
		lua_rawset(L, LUA_REGISTRYINDEX);
	}

	static int transform_ai_action(lua_State* L, ai::action_result_ptr action_result)
	{
		lua_newtable(L);
		lua_pushboolean(L, action_result->is_ok());
		lua_setfield(L, -2, "ok");
		lua_pushboolean(L, action_result->is_gamestate_changed());
		lua_setfield(L, -2, "gamestate_changed");
		lua_pushinteger(L, action_result->get_status());
		lua_setfield(L, -2, "status");
		return 1;
	}

	static bool to_map_location(lua_State* L, int& index, map_location& res)
	{
		if (lua_isuserdata(L, index))
		{
			uint const *u = luaW_touint(L, index);
			if (!u)
				return false;
			res = u->get_location();
			++index;
		}
		else
		{
			if (!lua_isnumber(L, index)) return false;
			res.x = lua_tointeger(L, index) - 1;
			++index;
			if (!lua_isnumber(L, index)) return false;
			res.y = lua_tointeger(L, index) - 1;
			++index;
		}

		return true;
	}

	static int ai_execute_move(lua_State* L, bool remove_movement)
	{
		int index = 1;
		if (false)
		{
error_call_destructors:
			return luaL_typerror(L, index, "location (unit/integers)");
		}

		int side = ((ai::engine_lua*)lua_touserdata(L, lua_upvalueindex(1)))->get_readonly_context().get_side();
		map_location from, to;
		if (!to_map_location(L, index, from)) goto error_call_destructions;
		if (!to_map_location(L, index, to)) goto error_call_destructions;
		ai::move_result_ptr move_result = ai::actions::execute_move_action(side, true, from, to, remove_movement);
		return transform_ai_action(L, move_result);
	}

	static int cfun_ai_execute_move_full(lua_State* L)
	{
		return ai_execute_move(L, true);
	}

	static int cfun_ai_execute_move_partial(lua_State* L)
	{
		return ai_execute_move(L, false);
	}

	static int cfun_ai_execute_attack(lua_State* L)
	{
		int index = 1;
		if (false)
		{
error_call_destructors:
			return luaL_typerror(L, index, "location (unit/integers)");
		}

		ai::readonly_context &context = ((ai::engine_lua*)lua_touserdata(L, lua_upvalueindex(1)))->get_readonly_context();

		int side = context.get_side();
		map_location attacker, defender;
		if (!to_map_location(L, index, attacker)) goto error_call_destructors;
		if (!to_map_location(L, index, defender)) goto error_call_destructors;

		int attacker_weapon = -1;
		double aggression = context.get_aggression();

		if (!lua_isnoneornil(L, index + 1) && lua_isnumber(L, index + 1))
		{
			aggression = lua_tonumber(L, index + 1);
		}

		if (!lua_isnoneornil(L, index))
		{
			attacker_weapon = lua_tointeger(L, index);
		}

		ai::attack_result_ptr attack_result = ai::actions::execute_attack_action(side, true, attacker, defender, attack_weapon, aggression);
		return transform_ai_action(L, attack_result);
	}

	static int ai_execute_stopunit_select(lua_State* L, bool remove_movement, bool remove_attacks)
	{
		int index = 1;
		if (false)
		{
error_call_destructors:
			return luaL_typerror(L, index, "location (unit/integers)");
		}

		int side = ((ai::engine_lua*)lua_touserdata(L, lua_upvalueindex(1)))->get_readonly_context().get_side();
		map_location loc;
		if (!to_map_location(L, index, loc)) goto error_call_destructors;

		ai::stopunit_result_ptr stopunit_result = ai::actions::execute_stopunit_action(side, true, loc, remove_movement, remove_attacks);
		return transform_ai_action(L, stopunit_result);
	}

	static int cfun_ai_execute_stopunit_moves(lua_State* L)
	{
		return ai_execute_stopunit_select(L, true, false);
	}

	static int cfun_ai_execute_stopunit_attacks(lua_State* L)
	{
		return ai_execute_stopunit_select(L, false, true);
	}

	static int cfun_ai_execute_stopunit_all(lua_State* L)
	{
		return ai_execute_stopunit_select(L, true, true);
	}

	static int cfun_ai_execute_recruit(lua_State* L)
	{
		const char* unit_name = luaL_checkstring(L, 1);
		int side = ((ai::engine_lua*)lua_touserdata(L, lua_upvalueindex(1)))->get_readonly_context().get_side();
		map_location where;
		if (!lua_isnoneornil(L, 2))
		{
			where.x = lua_tonumber(L, 2) - 1;
			where.y = lua_tonumber(L, 3) - 1;
		}

		ai::recruit_result_ptr recruit_result = ai::actions::execute_recruit_action(side, true, std::string(unit_name), where);
		return transform_ai_action(L, recruit_result);
	}

	static int cfun_ai_execute_recall(lua_State* L)
	{
		const char* unit_id = luaL_checkstring(L, 1);
		int side = ((ai::engine_lua*)lua_touserdata(L, lua_upvalueindex(1)))->get_readonly_context().get_side();
		map_location where;
		if (!lua_isnoneornil(L, 2))
		{
			where.x = lua_tonumber(L, 2) - 1;
			where.y = lua_tonumber(L, 3) - 1;
		}

		ai::recall_result_ptr recall_result = ai::actions::execute_recall_action(side, true, std::string(unit_id), where);
		return transform_ai_action(L, recall_result);
	}

	lua_ai_context* lua_ai_context::create(lua_State* L, char const* code, ai::engine_lua* engine)
	{
		int res_ai = luaL_loadstring(L, code);
		if (res_ai)
		{
			char const *m = lua_tostring(L, -1);
			ERR_LUA << "error while initializing ai: " << m << '\n';
			lua_pop(L, 2);
			return NULL;
		}

		lua_newtable(L);
		lua_pushinteger(L, engine->get_readonly_context().get_side());

		lua_setfield(L, -2, "side");
		
		static luaL_reg const callbacks[] = 
		{
			{"attack", &cfun_ai_execute_attack},
			{"move", &cfun_ai_execute_move_partial},
			{"move_full", &cfun_ai_execute_move_full}, 
			{"recall", &cfun_ai_execute_recall},
			{"recruit", &cfun_ai_execute_recruit},
			{"stopunit_all", &cfun_ai_execute_stopunit_all},
			{"stopunit_attacks", &cfun_ai_execute_stopunit_attacks},
			{"stopunit_moves", &cfun_ai_execute_stopunit_moves},
			{NULL, NULL}
		};

		for (const luaL_reg* p = callbacks; p->name; ++p)
		{
			lua_pushlightuserdata(L, engine);
			lua_pushcclosure(L, p->func, 1);
			lua_setfield(L, -2, p->name);
		}

		if (!luaW_pcall(L, 1, 1, true))
		{
			return NULL;
		}

		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);

		size_t length_ai = lua_objlen(L, -1);
		lua_pushvalue(L, -2);
		lua_rawseti(L, -2, length_ai + 1);
		lua_pop(L, 2);
		return new lua_ai_context(L, length_ai + 1, engine->get_readonly_context().get_side());
	}

	lua_ai_action_handler* lua_ai_action_handler::create(lua_State* L, char const* code, lua_ai_context& context)
	{
		int res = luaL_loadstring(L, code);
		if (res)
		{
			char const* m = lua_tostring(L, -1);
			ERR_LUA << "error while creating ai function: " << m << '\n';
			lua_pop(L, 2);
			return NULL;
		}

		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);
		size_t length = lua_objlen(L, -1);
		lua_pushvalue(L, -2);
		lua_rawseti(L, -2, length + 1);
		lua_remove(L, -1);
		lua_remove(L, -1);

		return new lua_ai_action_handler(L, context, length + 1);
	}

	void lua_ai_context::load()
	{
		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);
		lua_rawgeti(L, -1, num_);
		lua_remove(L, -2);
	}

	lua_ai_context::~lua_ai_context()
	{
		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);
		lua_pushnil(L);
		lua_rawseti(L, -2, num_);
		lua_pop(L, 1);
	}

	void lua_ai_action_handler::handle(config& cfg, bool configOut)
	{
		int initial_top = lua_gettop(L);

		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);
		lua_rawgeti(L, -1, num_);
		lua_remove(L, -2);
		context_.load();

		if (!configOut)
		{
			luaW_pushconfig(L, cfg);
			luaW_pcall(L, 2, 0, true);
		}
		else if (luaW_pcall(L, 1, 2, true))
		{
			int score = lua_tointeger(L, initial_top + 1);
			if (!luaW_toconfig(L, initial_top + 2, cfg))
			{
				ERR_LUA << "Failed to parsed WML_object returned by AI function.\n"
			}
			cfg["score"] = score;
		}
		lua_settop(L, initial_top);
	}

	lua_ai_action_handler::~lua_ai_action_handler()
	{
		lua_pushlightuserdata(L, (void*)&aisKey);
		lua_rawget(L, LUA_REGISTRYINDEX);
		lua_pushnil(L);
		lua_rawseti(L, -2, num_);
		lua_pop(L, 1);
	}
}
